"""
模块名：hxpTTS
功能：调用 GPT-SoVITS 模型实现文本到语音的推理
作者：hxp@海小皮
版本： V1.1
依赖：
    DataStruct
    TextProcess
    torch psutil os json time
    numpy librosa
    cnhubert BertModel GPTModel SoVITSModel
    sounddevice
"""
from .DataStruct import Emotion, EmotionType, GameType, VersionType, DeviceType
import PlotDub.TextProcess as tp
from typing import Dict
import torch, json
import psutil, os
from time import time as ttime
import numpy as np
import librosa
from feature_extractor import cnhubert
cnhubert_base_path = "GPT_SoVITS/pretrained_models/chinese-hubert-base"
cnhubert.cnhubert_base_path = cnhubert_base_path
from .BertModel import Bert
from .GPTModel import GPT
from .SoVITSModel import SoVITS
from collections import OrderedDict


class InferenceParam:
    """
    功能：推理参数数据结构
    成员：
        language 推理语言
        version 调用的 GPT-SoVITS 模型的版本
        if_freeze 默认值 False
        speed 语速 默认值 1
        top_k 默认值 15
        top_p 默认值 1
        temperature 默认值 1
    注意：
        language 和 version 参数在初始化时给定，加载文件函数不会修改这两个属性
    """
    def __init__(self, language="中文", version=VersionType.V2, if_freeze=False, speed=1, top_k=15, top_p=1, temperature=1):
        """
        language 推理语言
        version 调用的 GPT-SoVITS 模型的版本 VersionType 枚举类型
        if_freeze 默认值 False
        speed 语速 默认值 1
        top_k 默认值 15
        top_p 默认值 1
        temperature 默认值 1
        """
        self.language = language
        self.version = version
        self.if_freeze = if_freeze
        self.speed = speed
        self.top_k = top_k
        self.top_p = top_p
        self.temperature = temperature

    def LoadFromJson(self, path):
        """
        功能：从 json 文件中加载 InferenceParam 参数
        输入：
            path json 文件路径
        注意：
            不会修改 language 和 version 属性
        """
        with open(f"{path}", 'r', encoding='utf-8') as file:
            inference_param_data = json.load(file)
            self.if_freeze = inference_param_data.get("if_freeze", bool)
            self.speed = inference_param_data.get("speed")
            self.top_k = inference_param_data.get("top_k")
            self.top_p = inference_param_data.get("top_p")
            self.temperature = inference_param_data.get("temperature")

    def __repr__(self):
        return f"【推理参数InferenceParam】(【推理语言】={self.language}, 【GPT-SoVITS模型版本】={self.version}, 【if_freeze】={self.if_freeze}, 【speed】={self.speed}, 【top_k】={self.top_k}, 【top_p】={self.top_p}, 【temperature】={self.temperature})"


class Role:
    """
    功能：游戏角色信息类型
    成员：
        game 角色所属游戏 GameType 枚举类型
        name 角色名
        gptPath GPT 模型权重文件路径
        sovitsPath SoVITS 模型权重文件路径
        emotionsDict 角色情感信息字典 key = EmotionType value = Emotion 对象
    """
    def __init__(self, game: GameType, name, gpt_path, sovits_path, emo_dict: Dict[EmotionType, Emotion]):
        """
        game 角色所属游戏 GameType 枚举类型
        name 角色名字
        gpt_path GPT 模型权重文件路径
        sovits_path SoVITS 模型权重文件路径
        emo_dict 情感数据字典 key = EmotionType value = Emotion 对象
        """
        self.game = game
        self.name = name
        self.gptPath = gpt_path
        self.sovitsPath = sovits_path
        self.emotionsDict = emo_dict

    def __repr__(self):
        return f"【角色信息Role】\n(【游戏名】={self.game}, 【角色名】={self.name}, \n【gpt路径】={self.gptPath}, \n【sovits路径】={self.sovitsPath}, \n【情感信息】={self.emotionsDict})"


class ModelParam:
    """
    功能：模型调用相关参数及参数加载
    成员：
        device 硬件类型 DeviceType 枚举类型
        gpuNums GPU 数量
        isHalf 是否可启用半精度
    """
    def __init__(self, version=VersionType.V2):
        self.device = DeviceType.CPU
        self.gpuNums = 0
        self.isHalf = False
        self.__initModuleParam()
        # 定义推理语言字典
        self.languagesDict = {
            "中文": "all_zh",  # 全部按中文识别
            "英文": "en",  # 全部按英文识别#######不变
            "日文": "all_ja",  # 全部按日文识别
            "中英混合": "zh",  # 按中英混合识别####不变
            "日英混合": "ja",  # 按日英混合识别####不变
            "多语种混合": "auto",  # 多语种启动切分识别语种
        } if version == VersionType.V1 else {
            "中文": "all_zh",  # 全部按中文识别
            "英文": "en",  # 全部按英文识别#######不变
            "日文": "all_ja",  # 全部按日文识别
            "粤语": "all_yue",  # 全部按中文识别
            "韩文": "all_ko",  # 全部按韩文识别
            "中英混合": "zh",  # 按中英混合识别####不变
            "日英混合": "ja",  # 按日英混合识别####不变
            "粤英混合": "yue",  # 按粤英混合识别####不变
            "韩英混合": "ko",  # 按韩英混合识别####不变
            "多语种混合": "auto",  # 多语种启动切分识别语种
            "多语种混合(粤语)": "auto_yue",  # 多语种启动切分识别语种
        }

    def __initModuleParam(self):
        # 用于判断有效的 N 卡
        ok_gpu_keywords = {"10", "16", "20", "30", "40", "A2", "A3", "A4", "P4", "A50", "500", "A60", "70", "80", "90",
                           "M4", "T4", "TITAN", "L4", "4060", "H"}
        # 设置 PyTorch 的随机种子，使得实验结果具有可重复性
        torch.manual_seed(233333)
        # 识别显卡
        gpu_numbers_set = set()
        gpu_infos = []
        mem = []
        gpu_ok = False
        if torch.cuda.is_available():
            self.device = DeviceType.GPU
            self.gpuNums = torch.cuda.device_count()
            self.isHalf = True
            for i in range(self.gpuNums):
                gpu_name = torch.cuda.get_device_name(i)
                if any(value in gpu_name.upper() for value in ok_gpu_keywords):
                    # A10#A100#V100#A40#P40#M40#K80#A4500
                    gpu_ok = True  # 至少有一张能用的N卡
                    gpu_infos.append("%s\t%s" % (i, gpu_name))
                    gpu_numbers_set.add(i)
                    mem.append(int(torch.cuda.get_device_properties(i).total_memory / 1024 / 1024 / 1024 + 0.4))
            if gpu_ok and len(gpu_infos) > 0:
                gpu_info = "\n".join(gpu_infos)
                default_batch_size = min(mem) // 2
            else:
                gpu_info = ("%s\t%s" % ("0", "CPU"))
                gpu_infos.append("%s\t%s" % ("0", "CPU"))
                gpu_numbers_set.add(0)
                default_batch_size = int(psutil.virtual_memory().total / 1024 / 1024 / 1024 / 2)
            gpus = "-".join([i[0] for i in gpu_infos])

            default_gpu_numbers = str(sorted(list(gpu_numbers_set))[0])

            # 将越界的 number 强制改到界内
            def fix_gpu_number(input):
                try:
                    if (int(input) not in gpu_numbers_set): return default_gpu_numbers
                except:
                    return input
                return input

            os.environ["CUDA_VISIBLE_DEVICES"] = fix_gpu_number(gpus)
        else:
            self.device = DeviceType.CPU


class hxpTTS:
    """
    功能：Text To Speach 文本语音合成
    成员：
        __gamesRolesDict 游戏角色字典 key = game+Role.name value = Role 对象
        __inferenceParam 模型推理参数
        __cache 存放推理缓存
        __modelParam 模型加载相关参数
        __sslModel ssl 模型对象
        __bertModel bert 模型对象
        __GPTModel GPT 模型对象
        __SoVITSModel SoVITS 模型对象
    """
    def __init__(self, games_roles_dict: Dict[str, Role], inference_param: InferenceParam):
        """
        功能：初始化 TTS 对象，包括推理参数，权重字典，加载前置模型
        输入：
            game_roles_dict 游戏角色模型及情感参数字典 key = Role.game.value+Role.name value = Role 对象
            inference_param InferenceParam 结构模型推理参数
        """
        self.__gamesRolesDict = games_roles_dict
        self.__inferenceParam = inference_param
        self.__cache = OrderedDict()
        self.__cache_max_size = 10  # 硬编码缓存区域的大小
        # 初始化 torchInform
        self.__modelParam = ModelParam(inference_param.version)
        # 加载 ssl 模型
        self.__sslModel = cnhubert.get_model()
        if self.__modelParam.isHalf:
            self.__sslModel = self.__sslModel.half().to(self.__modelParam.device.value)
        else:
            self.__sslModel = self.__sslModel.to(self.__modelParam.device.value)
        # 加载 bert 模型
        self.__bertModel = Bert(self.__modelParam.isHalf, self.__modelParam.device)
        # 加载 GPT 模型
        self.__GPTModel = GPT(self.__modelParam.isHalf, self.__modelParam.device)
        # 加载 SoVITS 模型
        self.__SoVITSModel = SoVITS(self.__modelParam.languagesDict, self.__modelParam.isHalf, self.__modelParam.device)

    def AddGameRole(self, role: Role):
        """
        功能：添加游戏角色信息到 __gameRolesDict，若存在则刷新
        输入：role Role 结构对象
        """
        self.__gamesRolesDict[role.game.value + role.name] = role
        print(f"【{self.__class__.__name__}】【{self.AddGameRole.__name__}】添加【{role.game}】角色【{role.name}】")

    def SetParam(self, inference_param: InferenceParam):
        """
        功能：更新模型推理参数 __inferenceParam（仅推理部分，不涉及模型版本和推理语言）
        输入：inference_param InferenceParam 结构对象
        """
        self.__inferenceParam.if_freeze = inference_param.if_freeze
        self.__inferenceParam.speed = inference_param.speed
        self.__inferenceParam.top_p = inference_param.top_p
        self.__inferenceParam.top_k = inference_param.top_k
        self.__inferenceParam.temperature = inference_param.temperature
        print(f"【{self.__class__.__name__}】【{self.SetParam.__name__}】新的推理参数设置成功")

    def Inference(self, gpt_path, sovits_path, ref_wav, ref_text, tar_text):
        """
        功能：模型调用（仅支参考文本和目标文本同语种情况）
        输入：
            gpt_path GPT 模型权重路径
            sovits_path SoVITS 模型权重路径
            ref_wav 参考音频文件路径
            ref_text 参考文本
            tar_text 要合成的目标文本
        输出：
            合成的音频
        """
        # 检查输入合法性
        if len(gpt_path) == 0:
            print(f"【{self.__class__.__name__}】【{self.Inference.__name__}】没有有效的 GPT 模型权重路径")
        if len(sovits_path) == 0:
            print(f"【{self.__class__.__name__}】【{self.Inference.__name__}】没有有效的 SoVITS 模型权重路径")
        if len(ref_wav) == 0:
            print(f"【{self.__class__.__name__}】【{self.Inference.__name__}】没有有效的参考音频文件路径")
        if len(ref_text) == 0:
            print(f"【{self.__class__.__name__}】【{self.Inference.__name__}】没有有效的参考文本")
        if len(tar_text) == 0:
            print(f"【{self.__class__.__name__}】【{self.Inference.__name__}】没有有效的目标推理文本")

        # 更新模型参数
        self.__GPTModel.change_gpt_weights(gpt_path)
        self.__SoVITSModel.change_sovits_weights(sovits_path)

        # 时间记录列表，用于性能统计
        t = []

        # 如果未提供参考文本，则进入参考自由模式
        ref_free = False
        if ref_text is None or len(ref_text) == 0:
            ref_free = True

        # 记录起始时间
        t0 = ttime()

        # 将语言映射为模型支持的格式
        prompt_language = self.__modelParam.languagesDict[self.__inferenceParam.language]
        text_language = self.__modelParam.languagesDict[self.__inferenceParam.language]

        # 如果不在参考自由模式下，加载并处理参考音频
        if not ref_free:
            ref_text = ref_text.strip("\n")
            if (ref_text[-1] not in tp.splits): ref_text += "。" if prompt_language != "en" else "."
            print(f"【{self.__class__.__name__}】【{self.Inference.__name__}】实际输入的参考文本:{ref_text}")

        text = tar_text.strip("\n")  # 去除目标文本的换行符
        print(f"【{self.__class__.__name__}】【{self.Inference.__name__}】实际输入的目标文本:{text}")

        zero_wav = np.zeros(
            int(self.__SoVITSModel.hps.data.sampling_rate * 0.3),
            dtype=np.float16 if self.__modelParam.isHalf == True else np.float32,
        )

        # 加载参考音频并检查长度
        if not ref_free:
            with torch.no_grad():
                wav16k, sr = librosa.load(ref_wav, sr=16000)
                if (wav16k.shape[0] > 160000 or wav16k.shape[0] < 48000):
                    print(f"【{self.__class__.__name__}】【{self.Inference.__name__}】参考音频在3~10秒范围外，请更换！")
                    raise OSError("参考音频在3~10秒范围外，请更换！")

                # 转换为 Torch Tensor 并添加额外的静音部分
                device = str(self.__modelParam.device.value)
                wav16k = torch.from_numpy(wav16k).to(device)
                zero_wav_torch = torch.from_numpy(zero_wav)
                if self.__modelParam.isHalf:
                    wav16k = wav16k.half().to(device)
                    zero_wav_torch = zero_wav_torch.half().to(device)
                else:
                    wav16k = wav16k.to(device)
                    zero_wav_torch = zero_wav_torch.to(device)
                wav16k = torch.cat([wav16k, zero_wav_torch])

                # 提取 SSL 模型的语义特征
                ssl_content = self.__sslModel.model(wav16k.unsqueeze(0))[
                    "last_hidden_state"
                ].transpose(
                    1, 2
                )
                codes = self.__SoVITSModel.vq_model.extract_latent(ssl_content)
                prompt_semantic = codes[0, 0]
                prompt = prompt_semantic.unsqueeze(0).to(device)

        t1 = ttime()  # 时间记录
        t.append(t1 - t0)  # 记录第一阶段耗时

        # 提取参考文本的特征
        if not ref_free:
            phones1, bert1, norm_text1 = self.__bertModel.get_phones_and_bert(ref_text, prompt_language, self.__inferenceParam.version)

        # 输出当前合成的目标语句
        if (text[-1] not in tp.splits): text += "。" if text_language != "en" else "."
        print(f"【{self.__class__.__name__}】【{self.Inference.__name__}】开始合成语音:", text)

        # 提取目标文本的特征
        phones2, bert2, norm_text2 = self.__bertModel.get_phones_and_bert(text, text_language, self.__inferenceParam.version)
        print(f"【{self.__class__.__name__}】【{self.Inference.__name__}】bert提取的文本语义:", norm_text2)

        # 根据是否参考自由模式决定如何合并参考和目标特征
        if not ref_free:
            bert = torch.cat([bert1, bert2], 1)
            all_phoneme_ids = torch.LongTensor(phones1 + phones2).to(device).unsqueeze(0)
        else:
            bert = bert2
            all_phoneme_ids = torch.LongTensor(phones2).to(device).unsqueeze(0)

        bert = bert.to(device).unsqueeze(0)
        all_phoneme_len = torch.tensor([all_phoneme_ids.shape[-1]]).to(device)

        t2 = ttime()  # 记录推理开始时间

        # 检查缓存中是否已有语义特征
        if (text in self.__cache and self.__inferenceParam.if_freeze == True):
            pred_semantic = self.__cache[text]
        else:
            # 调用语义生成模型
            with torch.no_grad():
                pred_semantic, idx = self.__GPTModel.t2s_model.model.infer_panel(
                    all_phoneme_ids,
                    all_phoneme_len,
                    None if ref_free else prompt,
                    bert,
                    # prompt_phone_len=ph_offset,
                    top_k=self.__inferenceParam.top_k,
                    top_p=self.__inferenceParam.top_p,
                    temperature=self.__inferenceParam.temperature,
                    early_stop_num=self.__GPTModel.hz * self.__GPTModel.max_sec,
                )
                pred_semantic = pred_semantic[:, -idx:].unsqueeze(0)
                self.__cache[text] = pred_semantic  # 缓存语义特征
                # 如果超过阈值，删除最早插入的元素
                if len(self.__cache) > self.__cache_max_size:
                    self.__cache.popitem(last=False)  # last=False 删除第一个插入的元素

        t3 = ttime()  # 推理结束时间

        # 加载参考音频或额外参考音频用于解码
        refers = [self.__SoVITSModel.get_spepc(ref_wav).to(self.__bertModel.dtype).to(device)]

        # 解码生成音频
        audio = self.__SoVITSModel.vq_model.decode(
            pred_semantic,
            torch.LongTensor(phones2).to(device).unsqueeze(0),
            refers,
            speed=self.__inferenceParam.speed
        ).detach().cpu().numpy()[0, 0]

        # 防止爆音
        max_audio = np.abs(audio).max()  # 简单防止16bit爆音
        if max_audio > 1: audio /= max_audio

        # 添加静音部分
        audio_opt = [audio, zero_wav]

        t4 = ttime()  # 解码结束时间
        t.extend([t2 - t1, t3 - t2, t4 - t3])

        # 打印时间统计
        # print("%.3f\t%.3f\t%.3f\t%.3f" %
        #       (t[0], sum(t[1::3]), sum(t[2::3]), sum(t[3::3]))
        #       )

        # 输出最终音频
        return self.__SoVITSModel.hps.data.sampling_rate, (np.concatenate(audio_opt, 0) * 32768).astype(np.int16)

    def TTS(self, game: GameType, role_name, emo: EmotionType, text: str):
        """
        功能，根据输入 游戏 + 角色 + 情感 直接做单次推理
        输入：
            game GameType 枚举类型游戏名
            role_name 角色名
            emotion EmotionType 枚举类型情感
            text 要合成的文本
        输出：
            合成结果 合成成功 audio 合成失败 None
        """
        # 查找角色字典
        role = self.__gamesRolesDict.get(game.value + role_name)
        if role:
            print(f"【{self.__class__.__name__}】【{self.TTS.__name__}】【{game}】【{role_name}】【{emo.value}】【{text}】开始合成...")
        else:
            role = self.__gamesRolesDict.get(game.value + "默认")
            if role:
                print(f"【{self.__class__.__name__}】【{self.TTS.__name__}】【{game}】【{emo.value}】【{text}】使用默认音色进行合成...")
            else:
                print(f"【{self.__class__.__name__}】【{self.TTS.__name__}】【{game}】默认模型参数未找到！")
                return None

        # 解析模型参数
        gpt_path = role.gptPath
        sovits_path = role.sovitsPath
        ref_wav = role.emotionsDict[emo].refWav
        ref_text = role.emotionsDict[emo].refText
        # 推理调用
        audio = self.Inference(gpt_path, sovits_path, ref_wav, ref_text, text)

        return audio
